from typing import Sequence, Tuple, List
import multiprocessing as mp
import numpy as np
import time
import cv2
import os
import sys
import json
import apriltag

sys.path.append("../")
from arm_control.angle import ArmAngle
from arm_control.move import ArmMove

"""
读取保存为 JSON 格式的配置文件，并返回一个字典。
输入:
    file (str): 配置文件的路径。
输出：
    dict: 包含配置信息的字典。
"""
def load_config(file):
     # 检查文件是否存在
    if not os.path.isfile(file):
        raise FileNotFoundError(f"配置文件 {file} 不存在.")
     # 读取文件内容
    with open(file, 'r', encoding='utf-8') as f:
        config = json.load(f)
    return config

CUBE_HEIGHT = 3.8
MAX_NUMBER_OF_OBEJECTS = 4

PICK_AND_RELEASE_OFFSET = 6
ARM_Z_OFFSET = 5

WINDOW_W = 640
WINDOW_H = 480
WINDOW_NAME = "Project I: Apriltag Sort"

CONFIG_FILE = os.path.expanduser("~/.bc_config/cam_calibration.config")
cfg = load_config(CONFIG_FILE)

arm_base_coordinate_x = cfg["arm_base_coordinate"]["x"]
arm_base_coordinate_y = cfg["arm_base_coordinate"]["y"]
pixel_per_cm_ratio = cfg["pixel_per_cm_ratio"]

# 放置区域范围
place_position_in_pixels = {
    1: (arm_base_coordinate_x - int(10.5 * pixel_per_cm_ratio), arm_base_coordinate_y + int(13.5 * pixel_per_cm_ratio)),
    2: (arm_base_coordinate_x - int(10.5 * pixel_per_cm_ratio), arm_base_coordinate_y + int(20 * pixel_per_cm_ratio)),
    3: (arm_base_coordinate_x - int(10.5 * pixel_per_cm_ratio), arm_base_coordinate_y + int(26.5 * pixel_per_cm_ratio)),
    4: (arm_base_coordinate_x + int(10.5 * pixel_per_cm_ratio), arm_base_coordinate_y + int(13.5 * pixel_per_cm_ratio)),
    5: (arm_base_coordinate_x + int(10.5 * pixel_per_cm_ratio), arm_base_coordinate_y + int(20 * pixel_per_cm_ratio)),
    6: (arm_base_coordinate_x + int(10.5 * pixel_per_cm_ratio), arm_base_coordinate_y + int(26.5 * pixel_per_cm_ratio)),
}


def nothing(_):
    pass

"""
计算像素坐标系下的点距离原点的绝对距离
注意：这里的坐标系是以摄像机为原点,向右为x正方向,向下为y正方向方向.
输入:
    pixel_x (int): 像素坐标系下的x坐标
    pixel_y (int): 像素坐标系下的y坐标
    pixel_per_cm_ratio (float): 像素与厘米的比值, 1cm = ?pixel
输出:
    tuple (float, float): x方向和y方向上的距离(单位: cm)
"""
def compute_pixel_to_dist(
    pixel_x: int, pixel_y: int, pixel_per_cm_ratio: float
) -> Tuple[float, float]:
    dist_x, dist_y = pixel_x / pixel_per_cm_ratio, pixel_y / pixel_per_cm_ratio
    return round(dist_x, 2), round(dist_y, 2)

"""
给定起点和终点, 将其连线平均分为若干等份
输入:
    start_point (int, int): 像素坐标下的起始点
    end_point (int, int): 像素坐标下的终点
    num_points (int): 等分的份数
输出:
    list ((int, int)): num_points 个坐标, 每一个坐标用(x, y)表示
"""
def create_line_points(
    start_point: Tuple[int, int], end_point: Tuple[int, int], num_points: int = 6
) -> List[Tuple[int, int]]:
    points = []
    x_values = np.linspace(start_point[0], end_point[0], num_points)
    y_values = np.linspace(start_point[1], end_point[1], num_points)
    for x, y in zip(x_values, y_values):
        points.append((int(x), int(y)))
    return points

"""
两个点的欧式距离小于threshold,则认为他们重合，否则不重合.
输入:
    position_a tuple(int, int): 点a
    position_b tuple(int, int): 点b
输出:
    bool: True / False

"""
def is_matched(
    position_a: Tuple[int, int], position_b: Tuple[int, int], threshold: int = 5
) -> bool:
    if (position_a[0] - position_b[0]) ** 2 + (
        position_a[1] - position_b[1]
    ) ** 2 <= threshold**2:
        return True
    return False

"""
判断检测到的Apriltags 是否被放置到了目标位置。当检测有未被放置到目标位置的tags, 返回其id, x, y。如果没有,则返回-1, -1, -1
输入:
    target_positions list(points): 目标位置的列表，每一个元素为(x,y)
    tags list(tuple): 检测到的Apriltags 列表 (id, x, y)
输出:
    Tuple[int, int, int]: 当某个id的apriltag 位置和目标位置列表中对应index元素的坐标不重合时，返回[id, x, y], 如果没有的返回[-1, -1, -1]
"""
def mismatched_position(
    target_positions: List[Tuple[int, int]],
    tags: List[Tuple[int, int, int]],
    threshold: int = 5,
) -> Tuple[int, int, int]:
    for id, tag_x, tag_y in tags:
        res = is_matched(target_positions[id], (tag_x, tag_y), threshold)
        if not res:
            return id, tag_x, tag_y
    return -1, -1, -1

"""
在图像上绘制目标区域
输入:
    image (numpy array): 要绘制的图像
    points (list of tuple): 每个元素是一个二元组 (x, y)，表示一个点的坐标
    width (int): Default=40, 目标方块的边长
    radius (int): Default=4, 圆点的颜色, 默认为红色(BGR)
    alpha (float): default=0.3, 图像透明度
    color (tuple): 圆点的颜色，默认为绿色 (BGR格式)
    thickness (int): 线条的粗细。如果是负数（例如-1），则表示要绘制实心圆。
"""
def draw_target_pos(
    image: np.ndarray,
    points: List[Tuple[int, int]],
    width: int = 40,
    radius: float = 4,
    alpha: float = 0.3,
    color: Tuple[int, int, int] = (0, 255, 0),
    thickness: int = -1,
) -> np.ndarray:
    overlay = image.copy()
    for key, (x, y) in points.items():
        top_left = (int(x - width / 2), int(y - width / 2))
        bottom_right = (int(x + width / 2), int(y + width / 2))
        cv2.rectangle(overlay, top_left, bottom_right, color, -1)
        cv2.circle(overlay, (x, y), radius, (0, 0, 255), thickness)
        cv2.putText(
            overlay,
            f"ID:{key}",
            (top_left[0], top_left[1] - 10),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.5,
            color,
            2,
        )
    cv2.addWeighted(overlay, alpha, image, 1 - alpha, 0, image)
    return image


# 计算机械臂在摄像头像素坐标下距离原点的实际距离(x,y,z)
arm_base_pos = [
    *compute_pixel_to_dist(
        arm_base_coordinate_x, arm_base_coordinate_y, pixel_per_cm_ratio
    ),
    0,
]


def process_wrapper(queue, signal):
    def decorator(func):
        def wrapper():
            arm_move = ArmMove()
            arm_angle = ArmAngle(base_pos=arm_base_pos, z_offset=ARM_Z_OFFSET)
            while True:
                if not queue.empty():
                    targets = queue.get()
                    func(arm_move, arm_angle, *targets)
                    signal.set()
                time.sleep(0.1)  # Prevents busy waiting
        return wrapper
    return decorator


# Create a multiprocessing queue
queue = mp.Queue()
signal = mp.Event()
signal.set()

def _check_valid(step):
    """
    根据输入的step, 判断是否能够执行
    """
    if step is None:
        return False
    else:
        return True


"""
根据输入的目标点(x, y, z), 和操作类型（抓取或放置），
计算出机械臂的运动步进数，并执行相应的操作。
1. 先到达目标点上方offset处
2. 再下降到目标点, 并执行指定的操作（抓取或放置）
3. 再到达目标点上方offset处

位置坐标的基准为机械臂joint1的中心点

输入:
    arm_move[ArmMove]: 机械臂运动控制类的实例
    arm_angle[ArmAngle]: 机械臂动力学计算实例
    x[float]: 起始点x坐标, 单位: cm
    y[float]: 起始点y坐标, 单位: cm
    z[float]: 起始点z坐标, 单位: cm
    operation[str]: 'pick' 表示抓取操作，'release' 表示放置操作
    offset[float]: 上下偏移量, 单位: cm

输出:
    None
"""
def execute_operation(arm_move, arm_angle, x, y, z, operation="pick", offset=4):
    # Define a dictionary that maps operation names to functions.
    operations = {
        "pick": arm_move.pick,
        "release": arm_move.release,
        "go_to": arm_move.go_to,
        "reset": arm_move.reset,
    }

    if operation not in operations:
        print(f"无效的操作类型：{operation}")
        print(f"有效的操作类型为：{list(operations.keys())}")
        raise ValueError(f"无效的操作类型：{operation}\n有效的操作类型为:{list(operations.keys())}")

    if operation == "reset":
        print("\t1.回到初始位置")
        arm_move.go_to(0,0,0)
        return True


    # Calculate the steps once and reuse them in both operations.
    steps1 = arm_angle.get_steps(x, y, z + offset)
    steps2 = arm_angle.get_steps(x, y, z)

    if not all(_check_valid(step[0]) for step in [steps1, steps2]):
        return False

    print(f"\n\t当前任务为{operation}:")

    print(f"\t1.到达目标点 ({x, y, z}) 上方{offset}cm处。")
    arm_move.go_to(*steps1)
    time.sleep(0.2)

    if operation == "go_to":
        print("\t2.无需进一步操作，已到达目标点。")
        return True

    # Look up the appropriate function based on the operation parameter.
    operation_func = operations[operation]
    print(f"\t2.{'抓取' if operation == 'pick' else '放置'}目标物体")

    operation_func(*steps2)
    time.sleep(0.2)

    # Move to the hover position above the target point.
    print("\t3.位移到摆渡点")
    arm_move.go_to(*steps1)
    time.sleep(0.2)

    if operation == "release":
        print("\t4.回到初始位置")
        arm_move.go_to(0, 0, 0)
        print("\t当前任务完成\n")
    return True


"""
根据输入的(start_x, start_y, start_z)和(end_x, end_y, end_z),
计算出机械臂的运动步进数，并执行抓取操作。
1. 先到达目标点上方offset处
2. 再下降到目标点, 并执行抓取操作
3. 再到达目标点上方offset处
4. 再到达终点上方offset处
5. 再下降到终点, 并执行释放操作
6. 再到达终点上方offset处
7. 返回原点(0, 0, 0)

位置坐标的基准为机械臂joint1的中心点

输入:
    arm_move[ArmMove]: 机械臂运动控制类的实例
    arm_angle[ArmAngle]: 机械臂动力学计算实例
    start_x[float]: 起始点x坐标, 单位: cm
    start_y[float]: 起始点y坐标, 单位: cm
    start_z[float]: 起始点z坐标, 单位: cm
    end_x[float]: 终点x坐标, 单位: cm
    end_y[float]: 终点y坐标, 单位: cm
    end_z[float]: 终点z坐标, 单位: cm
    offset[float]: 上下偏移量, 单位: cm

输出:
    None
"""
def pick_and_release(
    arm_move,
    arm_angle,
    start_x,
    start_y,
    start_z,
    end_x,
    end_y,
    end_z,
    offset_pick=4,
    offset_release=4,
):
    operation_result = execute_operation(
        arm_move,
        arm_angle,
        start_x,
        start_y,
        start_z,
        operation="pick",
        offset=offset_pick,
    )

    if not operation_result:
        return

    operation_result = execute_operation(
        arm_move,
        arm_angle,
        end_x,
        end_y,
        end_z,
        operation="release",
        offset=offset_release,
    )

    if not operation_result:
        operation_result = execute_operation(
            arm_move,
            arm_angle,
            start_x,
            start_y,
            start_z,
            operation="release",
            offset=offset_pick,
        )


# Create a decorated version of pick_and_release
@process_wrapper(queue, signal)
def wrapped_pick_and_release(*args):
    pick_and_release(*args)

# tag检测，输入图像，反馈tag的数组
def tag_detect(self, frame: np.ndarray) -> List[Tuple[int, int, int]]:
    tags = []
    results = self.tag_detector.detect(frame)
    for result in results:
        tag_id = result.tag_id
        top_left, top_right, bottom_right, bottom_left = result.corners
        center_x = int((top_left[0] + bottom_right[0]) / 2)
        center_y = int((top_left[1] + bottom_right[1]) / 2)
        tag_inform = (tag_id, center_x, center_y)
        tags.append(tag_inform)
    return tags

# 在图像中画十字
def draw_cross(
    frame: np.ndarray,
    point: Sequence[int],
    size: int = 20,
    color: Tuple[int, int, int] = (0, 0, 255),
    thickness: int = 2,
) -> None:
    # 计算十字的四个端点坐标
    x, y = point
    p1 = (x - size, y)
    p2 = (x + size, y)
    p3 = (x, y - size)
    p4 = (x, y + size)
    # 在图像上画十字
    cv2.line(frame, p1, p2, color, thickness)
    cv2.line(frame, p3, p4, color, thickness)


"""
在图像上绘制检测到的apriltag的位置和id。
参数:
    frame (np.ndarray): 图像帧。
    id (int): 检测到的apriltag的id。
    point (Sequence[int]): 检测到的apriltag的位置，作为包含两个整数的序列 (x, y)。
    size (int, 可选): 标签的大小，默认值为20。
    color_cross (Tuple[int, int, int], 可选): 绘制apriltag十字准线的颜色，默认值为 (0, 0, 255) 表示红色。
    color_text (Tuple[int, int, int], 可选): 绘制apriltag的ID，默认值为 (255,255,0)。
    font_shift (Tuple[int, int], 可选): 调整绘制文字的位置偏移。
    thickness (int, 可选): 绘制标签的线条粗细，默认值为2。
返回:
     None
"""
def draw_tag(
    frame: np.ndarray,
    id: int,
    point: Tuple[int, int],
    size: int = 20,
    color_cross: Tuple[int, int, int] = (0, 0, 255),
    color_text: Tuple[int, int, int] = (255, 255, 0),
    font_shift: Tuple[int, int] = (25, 5),
    thickness: int = 2,
) -> None:
    draw_cross(frame, point, size, color_cross, thickness)
    cv2.putText(
        frame,
        f"ID:{id}",
        (point[0] + font_shift[0], point[1] + font_shift[1]),
        cv2.FONT_HERSHEY_SIMPLEX,
        0.8,
        color_text,
        1,
    )


def run(min_dist_threshold = 10):
    # Start the process
    process = mp.Process(target=wrapped_pick_and_release)
    process.start()
    time.sleep(3.0)
    
    cap = cv2.VideoCapture(0)
    detector = apriltag.Detector()
    
    cv2.namedWindow(WINDOW_NAME, cv2.WINDOW_NORMAL)
    cv2.resizeWindow(WINDOW_NAME, WINDOW_W, WINDOW_H)

    place_position_in_dist = {
        name: (*compute_pixel_to_dist(x, y, pixel_per_cm_ratio), CUBE_HEIGHT)
        for name, (x, y) in place_position_in_pixels.items()
    }
    target_box_size_in_pixel = int(CUBE_HEIGHT * pixel_per_cm_ratio)

    i = 0
    try:
        ### 整个while 是检测函数的主体, 在主体里面插入 pick_and_release 函数进程的控制信号量
        while True:
            ret, frame = cap.read()
            gray_img = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

            # 对输入的图像检测 Apriltag
            tag_pos_pairs = []
            results = detector.detect(gray_img)
            for result in results:
                tag_id = result.tag_id
                top_left, top_right, bottom_right, bottom_left = result.corners
                center_x = int((top_left[0] + bottom_right[0]) / 2)
                center_y = int((top_left[1] + bottom_right[1]) / 2)
                tag_inform = (tag_id, center_x, center_y)
                tag_pos_pairs.append(tag_inform)

            # 可视化检测结果
            for tag in tag_pos_pairs:
                draw_tag(frame, tag[0], tag[1:])

            # 可视化放置位置
            frame = draw_target_pos(frame, place_position_in_pixels, width=target_box_size_in_pixel)
            cv2.putText(
                frame,
                f"Jobs Count: {i}",
                (10, 30),
                cv2.FONT_HERSHEY_SIMPLEX,
                1,
                (0, 255, 0),
                2,
            )

            # 当前正在抓取中,不执行任何其它抓取判断
            if not signal.is_set():
                time.sleep(0.03)

            else:
                ##  控制信号量
                tag_pos_pairs_sorted_by_id = sorted(tag_pos_pairs, key=lambda x: x[0])

                # 判断检测到的tags 是否都在对应的位置, 返回第一个不对应的tag
                mismatched_id, x, y = mismatched_position(
                    place_position_in_pixels, tag_pos_pairs_sorted_by_id, min_dist_threshold
                )

                if mismatched_id != -1:
                    mismatched_apriltag_coordinate = compute_pixel_to_dist(
                        x, y, pixel_per_cm_ratio
                    )

                    print(f"将坐标发送到队列,  当前放置的Apriltag为 {mismatched_id}")
                    queue.put(
                        (
                            mismatched_apriltag_coordinate[0],
                            mismatched_apriltag_coordinate[1],
                            CUBE_HEIGHT,
                            place_position_in_dist[mismatched_id][0],
                            place_position_in_dist[mismatched_id][1],
                            place_position_in_dist[mismatched_id][2],
                            PICK_AND_RELEASE_OFFSET,
                            PICK_AND_RELEASE_OFFSET
                        )
                    )
                    i += 1
                    signal.clear()
            cv2.imshow(WINDOW_NAME,frame)

    except KeyboardInterrupt:
        print("Process interrupted by user")
    finally:
        process.terminate()
        process.join()


if __name__ == "__main__":
    min_dist_threshold = 15
    run(min_dist_threshold)
